// Copyright (c) Microsoft Corporation. All rights reserved.

using System;
using System.Diagnostics;
using System.Runtime.InteropServices;
using System.Runtime.Serialization;
using System.Text;
using System.Reflection;
using System.IO;
using System.Threading;

namespace Microsoft.Tools.CodeRunner
{
	/// <summary>
	/// The exception that is thrown when process injection fails
	/// </summary>
	[Serializable()]
	public sealed class ProcessInjectorException : SystemException
	{
		#region Public Instance Constructors

		/// <summary>
		/// Initializes a new instance of the <see cref="ProcessInjectorException" /> class.
		/// </summary>
		public ProcessInjectorException()
			: base()
		{
		}

		/// <summary>
		/// Initializes a new instance of the <see cref="ProcessInjectorException" /> class
		/// with a descriptive message.
		/// </summary>
		/// <param name="message">A descriptive message to include with the exception.</param>
		public ProcessInjectorException(string message)
			: base(message)
		{
		}

		/// <summary>
		/// Initializes a new instance of the <see cref="ProcessInjectorException" /> class
		/// with a descriptive message and an inner exception.
		/// </summary>
		/// <param name="message">A descriptive message to include with the exception.</param>
		/// <param name="innerException">A nested exception that is the cause of the current exception.</param>
		public ProcessInjectorException(string message, Exception innerException)
			: base(message, innerException)
		{
		}

		#endregion Public Instance Constructors

		#region Private Instance Constructors

		/// <summary>
		/// Initializes a new instance of the <see cref="ProcessInjectorException" /> class 
		/// with serialized data.
		/// </summary>
		/// <param name="info">The <see cref="SerializationInfo" /> that holds the serialized object data about the exception being thrown.</param>
		/// <param name="context">The <see cref="StreamingContext" /> that contains contextual information about the source or destination.</param>
		private ProcessInjectorException(SerializationInfo info, StreamingContext context)
			: base(info, context)
		{
		}

		#endregion Private Instance Constructors
	}

	/// <summary>
	/// A class for injecting a thread into a remote process.
	/// </summary>
	public unsafe class ProcessInjector : IDisposable
	{
		#region Native Methods
		/// <summary>
		/// See VirtualAlloc.
		/// </summary>
		[Flags]
		private enum AllocationType
		{
			/// <summary>
			/// See MEM_COMMIT
			/// </summary>
			Commit = 0x1000,
			/// <summary>
			/// See MEM_RESERVE
			/// </summary>
			Reserve = 0x2000,
			/// <summary>
			/// See MEM_DECOMMIT
			/// </summary>
			Decommit = 0x4000,
			/// <summary>
			/// See MEM_RELEASE
			/// </summary>
			Release = 0x8000,
			/// <summary>
			/// See MEM_RESET
			/// </summary>
			Reset = 0x80000,
			/// <summary>
			/// See MEM_PHYSICAL
			/// </summary>
			Physical = 0x400000,
			/// <summary>
			/// See MEM_TOP_DOWN
			/// </summary>
			TopDown = 0x100000
		}

		/// <summary>
		/// Memory page protection levels
		/// </summary>
		[Flags]
		private enum PageProtection
		{
			/// <summary>
			/// See PAGE_NOACCESS
			/// </summary>
			NoAccess = 0x01,
			/// <summary>
			/// Read only. See PAGE_READONLY
			/// </summary>
			ReadOnly = 0x02,
			/// <summary>
			/// Read and write.  See PAGE_READWRITE
			/// </summary>
			ReadWrite = 0x04,
			/// <summary>
			/// Copy on write.  See PAGE_WRITECOPY
			/// </summary>
			WriteCopy = 0x08,
			/// <summary>
			/// See PAGE_EXECUTE
			/// </summary>
			Execute = 0x10,
			/// <summary>
			/// See PAGE_EXECUTE_READ
			/// </summary>
			ExecuteRead = 0x20,
			/// <summary>
			/// See PAGE_EXECUTE_READWRITE
			/// </summary>
			ExecuteReadWrite = 0x40,
			/// <summary>
			/// See PAGE_EXECUTE_WRITECOPY
			/// </summary>
			ExecuteWriteCopy = 0x80,
			/// <summary>
			/// See PAGE_GUARD
			/// </summary>
			Guard = 0x100,
			/// <summary>
			/// See PAGE_NOCACHE
			/// </summary>
			NoCache = 0x200,
			/// <summary>
			/// See PAGE_WRITECOMBINE
			/// </summary>
			WriteCombine = 0x400
		}

		/// <summary>
		/// Process access flags
		/// </summary>
		[Flags]
		private enum ProcessAccessFlags
		{
			/// <summary>See PROCESS_TERMINATE</summary>
			Terminate = 0x0001,
			/// <summary>See PROCESS_CREATE_THREAD</summary>					    
			CreateThread = 0x0002,
			/// <summary>See PROCESS_SET_SESSIONID</summary>					     
			SetSessionId = 0x0004,
			/// <summary>See PROCESS_VM_OPERATION</summary>					         
			VMOperation = 0x0008,
			/// <summary>See PROCESS_VM_READ</summary>					   
			VMRead = 0x0010,
			/// <summary>See PROCESS_VM_WRITE</summary>					  
			VMWrite = 0x0020,
			/// <summary>See PROCESS_DUP_HANDLE</summary>					   
			DupHandle = 0x0040,
			/// <summary>See PROCESS_CREATE_PROCESS</summary>
			CreateProcess = 0x0080,
			/// <summary>See PROCESS_SET_QUOTA</summary>
			SetQuota = 0x0100,
			/// <summary>See PROCESS_SET_INFORMATION</summary>
			SetInformation = 0x0200,
			/// <summary>See PROCESS_QUERY_INFORMATION</summary>
			QueryInformation = 0x0400,
			/// <summary>See PROCESS_SUSPEND_RESUME</summary>
			SuspendResume = 0x0800,
			/// <summary>See PROCESS_ALL_ACCESS</summary>
			AllAccess = 0xFFF & 0x001F0000
		}

		/// <summary>
		/// Thread creation flags
		/// </summary>
		[Flags]
		private enum ThreadCreationFlags
		{
			/// <summary>
			/// No flags
			/// </summary>
			Nothing = 0x0,
			/// <summary>
			/// See CREATE_SUSPENDED
			/// </summary>
			CreateSuspended = 0x00000004
		}

		/// <summary>
		/// Special return values from WaitForSingleObject, etc.
		/// </summary>
		public class WaitReturnCode
		{
			/// <summary>
			/// Wait object zero is signalled
			/// </summary>
			public const uint Object0 = 0x00000000;
			/// <summary>
			/// Wait timed out
			/// </summary>
			public const uint Timeout = 0x00000102;
			/// <summary>
			/// The thread has not exited yet
			/// </summary>
			public const uint Pending = 0x00000103;
			/// <summary>
			/// Wait failed
			/// </summary>
			public const uint Failed = 0xFFFFFFFF;
			/// <summary>
			/// Wait abandoned for object 0
			/// </summary>
			public const uint Abandoned0 = 0x00000080;
		}

		private class NativeMethods
		{
			[DllImport("kernel32.dll", CharSet=CharSet.Auto, SetLastError=true, ExactSpelling=true)]
			public static extern bool CloseHandle(HandleRef href);

			[DllImport("kernel32.dll")]
			public static unsafe extern byte* VirtualAllocEx(IntPtr hProcess, byte* pAddress,
				UIntPtr size, AllocationType allocationType, PageProtection protect);			

			[DllImport("kernel32.dll")]
			public static unsafe extern bool VirtualFreeEx(IntPtr hProcess, byte* pAddress,
				UIntPtr size, AllocationType freeType);

			[DllImport("kernel32.dll", ExactSpelling=true)]
			public static extern IntPtr OpenProcess(ProcessAccessFlags desiredAccess, bool inheritHandle, int dwProcessId);

			[DllImport("kernel32.dll")]
			public static extern IntPtr GetModuleHandle(string moduleName);

			[DllImport("kernel32.dll", CharSet=CharSet.Ansi, ExactSpelling=true)]
			public static extern byte* GetProcAddress(IntPtr hModule, string procName);			

			[DllImport("kernel32.dll")]
			public static unsafe extern bool WriteProcessMemory(IntPtr hProcess, byte* pBaseAddress,
				byte* pBuffer, UIntPtr size, UIntPtr* pNumberOfBytesWritten);

			[DllImport("kernel32.dll")]
			public static unsafe extern bool WriteProcessMemory(IntPtr hProcess, byte* pBaseAddress,
				byte[] buffer, UIntPtr size, UIntPtr* pNumberOfBytesWritten);

			[DllImport("kernel32.dll", ExactSpelling=true)]
			public unsafe static extern IntPtr CreateRemoteThread(IntPtr hProcess,
				void* pSecurityAttributes, UIntPtr stackSize, byte* pStartAddress, 
				void* pParameter, ThreadCreationFlags creationFlags, int* pThreadId);

			[DllImport("kernel32.dll", SetLastError = true)]
			public unsafe static extern uint WaitForSingleObject(HandleRef hrHandle, int milliseconds);

			[DllImport("kernel32.dll", ExactSpelling=true)]
			public static extern bool GetExitCodeThread(IntPtr hThread, out int exitCode);
		}
		
		#endregion
        
		#region IDisposable and Related Members

		/// <summary>
		/// Destructor
		/// </summary>
		~ProcessInjector()
		{
			Dispose(false);
		}
		
		/// <summary>
		/// Dispose method
		/// </summary>
		public void Dispose()
		{
			Dispose(true);
			GC.SuppressFinalize(this);
		}
		
		/// <summary>
		/// Close method
		/// </summary>
		public void Close()
		{
			Dispose();
		}
		
		private void Dispose(bool disposing)
		{
			if (hThread != IntPtr.Zero)
			{
				NativeMethods.CloseHandle(new HandleRef(this, hThread));
				hThread = IntPtr.Zero;
			}

			if (hProcess != IntPtr.Zero)
			{
				// Check everything for valid values, in case the constructor croaked
				if (pThreadData != null)
				{
					NativeMethods.VirtualFreeEx(hProcess, pThreadData, UIntPtr.Zero, AllocationType.Release);
					pThreadData = null;
				}

				// Use of HandleRef here protects the previous call from inadvertant garbage collection
				NativeMethods.CloseHandle(new HandleRef(this, hProcess));
				hProcess = IntPtr.Zero;
			}
		}

		#endregion

		#region Public Instance Methods
		
		/// <summary>
		/// Start a thread in that process which loads an DLL in that process and calls an exported method in the DLL.
		/// </summary>
		/// <param name="processId">ID of the process to inject into</param>
		/// <param name="dllPath">The full or partial path to the DLL to inject</param>
		/// <param name="functionName"></param>
		public ProcessInjector(int processId, string dllPath, string functionName)
		{
			// Make sure that the DLL name is fully qualified, or who knows what we'll end up injecting
			dllFileName = new ParsedPath(dllPath, PathType.Automatic).MakeFullPath().ToString();

			if (dllFileName.Length > 255)
				throw new ArgumentException("DLL path is too long");
				
			if (functionName.Length > 255)
				throw new ArgumentException("Function name is too long");
			
			Process.EnterDebugMode();
			
			// Open the process
			hProcess = NativeMethods.OpenProcess(
				ProcessAccessFlags.CreateThread | ProcessAccessFlags.VMOperation | 
				ProcessAccessFlags.VMWrite | ProcessAccessFlags.VMRead, false, processId);

			if (hProcess == IntPtr.Zero)
			{
				throw new ProcessInjectorException("Unable to open process handle");
			}

			// Allocate memory for injected code and data.  Memory pointers in another process are in the address space of that 
			// process.  You must use Read/WriteProcessMemory to access them.  So go over to the other process and take a look!  
			// TopToBottom is really useful to watch where the new memory is allocated, or to examine the memory while debugging.
			pThreadData = NativeMethods.VirtualAllocEx(
				hProcess, null, (UIntPtr)((uint)sizeof(RemoteThreadData) + (uint)remoteThreadProc.Length), 
				AllocationType.Commit, PageProtection.ExecuteReadWrite);

			if (pThreadData == null)
				throw new ProcessInjectorException();

			// Create a remote thread structure and initialize it to zero
			RemoteThreadData threadData = new RemoteThreadData();

			// Set the guard bytes
			threadData.guard = 0xabababab;

			// GetProcAddress doesn't support Unicode, so do the conversion to ASCII in this process
			Encoding ascii = Encoding.ASCII;
			Encoding unicode = Encoding.Unicode;
			
			byte[] unicodeBytes = unicode.GetBytes(functionName);
			byte[] asciiBytes = Encoding.Convert(unicode, ascii, unicodeBytes);
				
			// Copy bytes from functionNameAscii to threadData structure
			Memory.Copy(&threadData.functionName, asciiBytes, asciiBytes.Length);
			
			// Copy the DLL name too
			Memory.Copy(&threadData.modulePath, dllFileName, dllFileName.Length + 1 /* for null terminator */ );
			
			IntPtr hKernel32 = NativeMethods.GetModuleHandle("kernel32.dll");

			if (hKernel32 == IntPtr.Zero)
				throw new ProcessInjectorException();

			// Get the addresses of the functions that we will use in the remote thread
			threadData.pfnLoadLibrary = NativeMethods.GetProcAddress(hKernel32, "LoadLibraryW");
			threadData.pfnFreeLibrary = NativeMethods.GetProcAddress(hKernel32, "FreeLibrary");
			threadData.pfnGetProcAddress = NativeMethods.GetProcAddress(hKernel32, "GetProcAddress");

			if (threadData.pfnLoadLibrary == null ||
				threadData.pfnFreeLibrary == null ||
				threadData.pfnGetProcAddress == null)
			{
				throw new ProcessInjectorException("Unable to get required kernel32 entry points");
			}

			// Write the remote thread data structure to the other process address space
			if (!NativeMethods.WriteProcessMemory(hProcess, pThreadData, 
				(byte*)&threadData, (UIntPtr)((uint)sizeof(RemoteThreadData)), null))
			{
				throw new ProcessInjectorException("Unable to write memory to remote process");
			}

			// Write the remote thread procedure to the other process address space
			if (!NativeMethods.WriteProcessMemory(hProcess, pThreadData + sizeof(RemoteThreadData),
				remoteThreadProc, (UIntPtr)((uint)remoteThreadProc.Length), null))
			{
				throw new ProcessInjectorException("Unable to write memory to remote process");
			}

			// TODO: Flush remote process CPU cache here using FlushInstructionCache...

			int threadId;

			hThread = NativeMethods.CreateRemoteThread(hProcess, null, UIntPtr.Zero, 
				pThreadData + sizeof(RemoteThreadData), pThreadData, ThreadCreationFlags.Nothing, &threadId);

			if (hThread == IntPtr.Zero)
			{
				throw new ProcessInjectorException("Unable to create remote thread");
			}

			Process.LeaveDebugMode();
		}
		
		/// <summary>
		/// Wait for the injected thread to exit.  If the thread never exits this function never returns.
		/// </summary>
		public void WaitForExit()
		{
			WaitForExit(Timeout.Infinite);
		}
		
		/// <summary>
		/// Wait for the injected thread to exit
		/// </summary>
		/// <param name="milliseconds">Number of milliseconds to wait</param>
		/// <returns><code>true</code> if the current instance receives a signal; otherwise, <code>false</code>.</returns>
		public bool WaitForExit(int milliseconds)
		{
			uint ret = NativeMethods.WaitForSingleObject(new HandleRef(this, hThread), milliseconds);	
			
			return (ret == WaitReturnCode.Object0);
		}
		
		#endregion

		#region Public Instance Properties

		private int exitCode;
		
		/// <summary>
		/// Description of property <see cref="ExitCode"/>.
		/// </summary>
		public int ExitCode
		{
		    get 
		    {
				if (hThread == IntPtr.Zero)
					throw new InvalidOperationException("Object has been disposed");
				
				bool ok = NativeMethods.GetExitCodeThread(hThread, out exitCode);
				
				if (!ok)
					throw new InvalidOperationException("Unable to get thread status");
				
				if (exitCode == WaitReturnCode.Pending)
					throw new InvalidOperationException("Thread still executing");
		    
		        return exitCode;
		    }
		}

		/// <summary>
		/// Gets the name of the DLL that was injected into the remote process.
		/// </summary>
		public string DllFileName
		{
		    get 
		    {
		        return dllFileName;
		    }
		}

		#endregion		
		
		#region  Private Instance Data
		
		// TODO: These should be SafeHandles
		private IntPtr hProcess = IntPtr.Zero;
		private IntPtr hThread = IntPtr.Zero;
		private byte* pThreadData = null;
		private string dllFileName = null;
		
		#endregion

		#region Remote Instance Data

		[StructLayout(LayoutKind.Explicit)]
		struct RemoteThreadData
		{
			[FieldOffset(0)]public byte* pfnLoadLibrary;		// LoadLibrary address
			[FieldOffset(4)] public byte* pfnFreeLibrary;		// FreeLibrary address
			[FieldOffset(8)] public byte* pfnGetProcAddress;	// GetProcAddress address
			[FieldOffset(12)] public char modulePath;		    // the DLL path
			[FieldOffset(524)] public byte functionName;	    // the called function name (ANSI characters)
			[FieldOffset(780)] public uint guard;				// guard bytes, also forces the structure size to be correct
			
			// The data for the injected function begins here when we set things up in memory. See below.
		}

		#endregion

		#region Private Class Data
		// This is just like the good old days when we used to define great gobs of 6502 and Z80 code in 
		// BASIC DATA statements.  This is a small blob of x86 code that is injected into
		// the remote process to act as a bootstrapper.  It simply loads the DLL specified in the remote
		// instance data above.  See the InjStub project for details of the code.
		private static byte[] remoteThreadProc =
		{
			0x53, 0x56, 0x8B, 0x74, 0x24, 0x04, 0x57, 0x8D,
			0x46, 0x0C, 0x50, 0x33, 0xDB, 0xFF, 0x16, 0x8B,
			0xF8, 0x85, 0xFF, 0x74, 0x17, 0x8D, 0x8E, 0x0C,
			0x02, 0x00, 0x00, 0x51, 0x57, 0xFF, 0x56, 0x08,
			0x85, 0xC0, 0x74, 0x04, 0xFF, 0xD0, 0x8B, 0xD8,
			0x57, 0xFF, 0x56, 0x04, 0x5F, 0x5E, 0x8B, 0xC3,
			0x5B, 0x90, 0xC2, 0x04, 0x00
		};
		
		#endregion
	}
}
